package middleware

import (
	"bytes"
	"context"
	"github.com/gin-gonic/gin"
	"io"
	"time"
)

type LogMiddlewareBuilder struct {
	allowReqBody  bool                                        // 是否打印请求体
	allowRespBody bool                                        // 是否打印响应体
	logFunc       func(context context.Context, al AccessLog) // 打印逻辑
}

func NewLogMiddlewareBuilder(logFunc func(context context.Context, al AccessLog)) *LogMiddlewareBuilder {
	return &LogMiddlewareBuilder{
		logFunc: logFunc,
	}
}

func (l *LogMiddlewareBuilder) AllowReqBody() *LogMiddlewareBuilder {
	l.allowReqBody = true
	return l
}

func (l *LogMiddlewareBuilder) AllowRespBody() *LogMiddlewareBuilder {
	l.allowRespBody = true
	return l
}

func (l *LogMiddlewareBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		// 1. 打印请求
		var accessLog AccessLog
		accessLog.url = ctx.Request.URL.String()
		accessLog.method = ctx.Request.Method
		// 究竟要不要打印请求体(可能很大)
		// 究竟是用info呢还是debug呢？
		if l.allowReqBody && ctx.Request.Body != nil {
			data, _ := ctx.GetRawData()
			// 需要将数据放回去
			ctx.Request.Body = io.NopCloser(bytes.NewBuffer(data))
			if len(data) > 1024 {
				accessLog.reqBody = string(data[:1024])
			} else {
				accessLog.reqBody = string(data[:len(data)])
			}
		}
		// 2. 打印响应
		if l.allowRespBody {
			// 替换ctx中的responseWrite
			ctx.Writer = &responseWriter{
				ResponseWriter: ctx.Writer,
				al:             &accessLog,
			}
		}
		start := time.Now()
		defer func() {
			accessLog.duration = time.Since(start)
			l.logFunc(ctx, accessLog)
		}()
		// 放行
		ctx.Next()
	}
}

// AccessLog 允许打印的日志对象
type AccessLog struct {
	statusCode int
	url        string
	method     string
	reqBody    string
	respBody   string
	duration   time.Duration
}

type responseWriter struct {
	gin.ResponseWriter
	al *AccessLog
}

func (w *responseWriter) Write(data []byte) (n int, err error) {
	w.al.respBody = string(data)
	return w.ResponseWriter.Write(data)
}

func (w *responseWriter) WriteHeader(code int) {
	w.al.statusCode = code
	w.ResponseWriter.WriteHeader(code)
}
